#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <signal.h>
#include <pthread.h>

#define BUFFER_SIZE 4096
#define LARGE_STRING_SIZE 2000
#define NUM_THREADS 4
#define NUM_REQUESTS 4

volatile sig_atomic_t running = 1;

typedef struct {
    char *ip;
    int port;
    int thread_id;
} ThreadArgs;

void signal_handler(int signum) {
    if (signum == SIGINT || signum == SIGTERM) {
        running = 0;
        printf("\nExiting...\n");
    }
}

void* client_thread(void* arg) {
    ThreadArgs* args = (ThreadArgs*)arg;
    int sock_fd;
    struct sockaddr_in server_addr;
    char buffer[BUFFER_SIZE];
    char command[BUFFER_SIZE];
    char large_string[LARGE_STRING_SIZE];

    // 生成测试用的大字符串
    memset(large_string, 'A', LARGE_STRING_SIZE - 1);
    large_string[LARGE_STRING_SIZE - 1] = '\0';

    // 创建socket
    sock_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (sock_fd == -1) {
        perror("Socket creation failed");
        return NULL;
    }

    // 设置服务器地址
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(args->port);
    if (inet_pton(AF_INET, args->ip, &server_addr.sin_addr) <= 0) {
        perror("Invalid address");
        return NULL;
    }

    // 连接服务器
    if (connect(sock_fd, (struct sockaddr*)&server_addr, sizeof(server_addr)) != 0) {
        perror("Connection failed");
        return NULL;
    }

    printf("Thread %d connected to server\n", args->thread_id);

    // 发送多次请求
    for (int i = 0; i < NUM_REQUESTS && running; i++) {
        // 轮流测试不同命令
        switch (i % 4) {
            case 0:
                snprintf(command, BUFFER_SIZE, "*3\r\n$3\r\nSET\r\n$4\r\nname\r\n$5\r\nalice\r\n");
                break;
            case 1:
                snprintf(command, BUFFER_SIZE, "*2\r\n$3\r\nGET\r\n$4\r\nname\r\n");
                break;
            case 2:
                snprintf(command, BUFFER_SIZE, "*1\r\n$4\r\nPING\r\n");
                break;
            case 3:
                snprintf(command, BUFFER_SIZE, "*2\r\n$4\r\nECHO\r\n$5\r\nhello\r\n");
                break;
        }

        // 发送命令
        if (send(sock_fd, command, strlen(command), 0) < 0) {
            perror("Send failed");
            break;
        }

        // 接收响应
        ssize_t n = recv(sock_fd, buffer, BUFFER_SIZE - 1, 0);
        if (n <= 0) {
            printf("Thread %d: Server disconnected\n", args->thread_id);
            break;
        }
        buffer[n] = '\0';
        printf("Thread %d, Request %d: Server response: %s", args->thread_id, i + 1, buffer);
    }

    close(sock_fd);
    printf("Thread %d finished\n", args->thread_id);
    return NULL;
}

int main(int argc, char *argv[]) {
    if (argc != 3) {
        printf("Usage: %s <ip> <port>\n", argv[0]);
        return 1;
    }

    pthread_t threads[NUM_THREADS];
    ThreadArgs thread_args[NUM_THREADS];

    // 设置信号处理
    signal(SIGINT, signal_handler);

    // 创建多个客户端线程
    for (int i = 0; i < NUM_THREADS; i++) {
        thread_args[i].ip = argv[1];
        thread_args[i].port = atoi(argv[2]);
        thread_args[i].thread_id = i + 1;

        if (pthread_create(&threads[i], NULL, client_thread, &thread_args[i]) != 0) {
            perror("Failed to create thread");
            return 1;
        }
    }

    // 等待所有线程完成
    for (int i = 0; i < NUM_THREADS; i++) {
        pthread_join(threads[i], NULL);
    }

    printf("All threads completed\n");
    return 0;
}